import pandas as pd
from libs.Objects import Objects
from libs.MySqlUtil import MySqlUtil
from libs.StockBase import StockBase

stockBase = StockBase()


class StockToFile(object):
    RESULT_FILE = "data/result.xlsx"

    def __init__(self):
        self.conn = MySqlUtil(**MySqlUtil.CONF)
        print(self.conn.get_dicts("SELECT VERSION()"))

    def if_exist(self, tradingDay):
        tradingDay = Objects.to_db_date(tradingDay)
        values = self.conn.get_values('select 1 from trading_emotion where day={day}'.format(day=tradingDay))
        return not len(values) == 0

    def delete_by_day(self, tradingDay):
        tradingDay = Objects.to_db_date(tradingDay)
        self.conn.execute('delete from trading_emotion where day={day}'.format(day=tradingDay))

    def to_db(self, rowData):
        if Objects.isEmpty(rowData):
            raise Exception("date is empty")
        sql = """insert into trading_emotion(day,red_count,green_count,limit_up,limit_down,break_down_codes,limit_up1_codes,limit_up2_codes,limit_up3_codes,limit_up_gt3_codes,amount)
         values (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) """
        self.conn.insert_batch(sql, [rowData])

    def query_to_df(self, endDay):
        print("query trading_emotion")
        startDay = Objects.to_db_date(Objects.getDay(-120))
        endDay = Objects.to_db_date(endDay)
        sql = ("select * from trading_emotion where day between {startDay} and {endDay} order by day".
               format(startDay=startDay, endDay=endDay))
        return pd.DataFrame(self.conn.get_dicts(sql))

    def convert_result(self, all_df=pd.DataFrame()):
        def convert_name(row, col):
            if Objects.isEmpty(row[col]):
                return ''
            else:
                arr = []
                for code in row[col].split('，'):
                    if len(code) == 9:
                        arr.append(stockBase.getStockName(code))
                    else:
                        arr.append(stockBase.getStockName(code[0:9]) + code[9:])
                return ",".join(arr)

        def clac_size(row, col):
            return 0 if Objects.isEmpty(row[col]) else len(row[col].split('，'))

        def calc_percent(row):
            all = row['limit_up']
            up1 = row['limit_up1']
            per_size = all - up1
            per = round(per_size * 100 / all)
            return "%s%%(%s)" % (str(per_size), str(per))

        """
        day	char(8)	YES	MUL	-	-
        red_count	int	YES	-	-	-
        green_count	int	YES	-	-	-
        limit_up	int	YES	-	-	-
        limit_down	int	YES	-	-	-
        break_down_codes	text	YES	-	-	-
        limit_up1_codes	text	YES	-	-	-
        limit_up2_codes	text	YES	-	-	-
        limit_up3_codes	text	YES	-	-	-
        limit_up_gt3_codes
        amount
        """
        # 'day', '上涨家数', '下跌家数', '涨停家数', '跌停家数','成交额(万亿)', '炸板家数', '首板数量', '首板以上占比', '二板数量', '三板数量', '三板以上家数',
        all_df['break_down'] = all_df.apply(clac_size, axis=1, args=('break_down_codes',))
        all_df['amount'] = all_df.apply(lambda x: round(x['amount'] / (1000 * 1000 * 1000 * 1000), 2), axis=1)
        all_df['limit_up1'] = all_df.apply(clac_size, axis=1, args=('limit_up1_codes',))
        all_df['limit_up1_percent'] = all_df.apply(calc_percent, axis=1)
        all_df['limit_up2'] = all_df.apply(clac_size, axis=1, args=('limit_up2_codes',))
        all_df['limit_up3'] = all_df.apply(clac_size, axis=1, args=('limit_up3_codes',))
        all_df['limit_up_gt3'] = all_df.apply(clac_size, axis=1, args=('limit_up_gt3_codes',))

        # '炸板明细', '首板明细', '二板明细', '三板明细', '三板以上明细'
        all_df['break_down_name'] = all_df.apply(convert_name, axis=1, args=('break_down_codes',))
        del all_df['break_down_codes']

        all_df['limit_up1_names'] = all_df.apply(convert_name, axis=1, args=('limit_up1_codes',))
        del all_df['limit_up1_codes']

        all_df['limit_up2_names'] = all_df.apply(convert_name, axis=1, args=('limit_up2_codes',))
        del all_df['limit_up2_codes']

        all_df['limit_up3_names'] = all_df.apply(convert_name, axis=1, args=('limit_up3_codes',))
        del all_df['limit_up3_codes']

        all_df['limit_up_gt3_names'] = all_df.apply(convert_name, axis=1, args=('limit_up_gt3_codes',))
        del all_df['limit_up_gt3_codes']

        all_df.columns = ['day', '上涨家数', '下跌家数', '涨停家数', '跌停家数', '成交额(万亿)', '炸板家数',
                          '首板数量', '首板以上占比', '二板数量', '三板数量', '三板以上家数',
                          '炸板明细', '首板明细', '二板明细', '三板明细', '三板以上明细']
        return all_df

    def appendFile(self, date, rowData):
        if Objects.isEmpty(date):
            raise Exception("date is empty")
        sheetName = self.getSheetName(date)

        prev_df = self.getPrevDf(date)

        if 'date' in prev_df.columns and not prev_df[(prev_df['date'] == date)].empty:
            print('The row of date %s exists! exit' % (date))
            return

        data = rowData
        new_df = pd.DataFrame(data)

        print("write xlsx sheet " + sheetName)
        with pd.ExcelWriter(self.RESULT_FILE, if_sheet_exists='overlay', mode='a') as writer:
            # 创建一个DataFrame对象
            all_df = pd.concat([prev_df, new_df])
            # 将DataFrame对象写入到Excel文件中
            all_df.to_excel(writer, sheet_name=sheetName, index=False, freeze_panes=(1, 0), header=True)

    def getPrevDf(self, date):
        sheetName = self.getSheetName(date)
        if self.existSheetName(date):
            exist_df = pd.read_excel(self.RESULT_FILE, sheet_name=sheetName)
            return exist_df
        else:
            return pd.DataFrame()

    def existSheetName(self, date):
        xlsx = pd.ExcelFile(self.RESULT_FILE)
        sheetName = self.getSheetName(date)
        xlsx.close()
        return sheetName in xlsx.sheet_names

    def getSheetName(self, date):
        return date[0:4] + '年'


if __name__ == '__main__':
    toFile = StockToFile()
    day = '2025-04-07'
    df = toFile.query_to_df(day)
    str_name = "data/%s_stock.xlsx" % (day)
    df = toFile.convert_result(df)
    Objects.to_excel(str_name, df, pd.DataFrame())
